use std::cmp::Ordering;
use std::mem;

/// A node in the AVL tree.
struct AVLTreeNode<K, V> {
    children: [Option<Box<AVLTreeNode<K, V>>>; 2],
    parent: Option<*mut AVLTreeNode<K, V>>,
    key: K,
    value: V,
    height: i32,
}

/// An AVL tree, a self-balancing binary search tree.
pub struct AVLTree<K, V, F> {
    root_node: Option<Box<AVLTreeNode<K, V>>>,
    compare_func: F,
    num_nodes: usize,
}

/// A comparison function for keys in the AVL tree.
pub type CompareFunc<K> = fn(&K, &K) -> Ordering;

impl<K, V> AVLTreeNode<K, V> {
    /// Creates a new node with the given key and value.
    fn new(key: K, value: V) -> Self {
        AVLTreeNode {
            children: [None, None],
            parent: None,
            key,
            value,
            height: 1,
        }
    }
}

impl<K, V, F> AVLTree<K, V, F>
where
    F: Fn(&K, &K) -> Ordering,
{
    /// Creates a new AVL tree with the given comparison function.
    pub fn new(compare_func: F) -> Self {
        AVLTree {
            root_node: None,
            compare_func,
            num_nodes: 0,
        }
    }

    /// Frees the subtree rooted at the given node.
    fn free_subtree(&mut self, node: Option<Box<AVLTreeNode<K, V>>>) {
        if let Some(mut node) = node {
            self.free_subtree(node.children[0].take());
            self.free_subtree(node.children[1].take());
        }
    }

    /// Frees the entire tree.
    pub fn free(&mut self) {
        self.free_subtree(self.root_node.take());
    }

    /// Returns the height of the subtree rooted at the given node.
    fn subtree_height(node: &Option<Box<AVLTreeNode<K, V>>>) -> i32 {
        node.as_ref().map_or(0, |n| n.height)
    }

    /// Updates the height of the given node based on its children's heights.
    fn update_height(&mut self, node: &mut AVLTreeNode<K, V>) {
        let left_height = Self::subtree_height(&node.children[0]);
        let right_height = Self::subtree_height(&node.children[1]);
        node.height = 1 + std::cmp::max(left_height, right_height);
    }

    /// Finds the side of the given node relative to its parent.
    fn node_parent_side(node: &AVLTreeNode<K, V>) -> usize {
        unsafe {
            if let Some(parent) = node.parent {
                if (*parent).children[0].as_ref().map_or(false, |n| n.key == node.key) {
                    0
                } else {
                    1
                }
            } else {
                panic!("Node has no parent");
            }
        }
    }

    /// Replaces node1 with node2 at its parent.
    fn node_replace(&mut self, node1: &mut AVLTreeNode<K, V>, node2: Option<Box<AVLTreeNode<K, V>>>) {
        if let Some(mut node2) = node2 {
            node2.parent = node1.parent;
        }

        if node1.parent.is_none() {
            self.root_node = node2;
        } else {
            let side = Self::node_parent_side(node1);
            unsafe {
                if let Some(parent) = node1.parent {
                    (*parent).children[side] = node2;
                    self.update_height(&mut *parent);
                }
            }
        }
    }

    /// Rotates the subtree rooted at the given node in the specified direction.
    fn rotate(&mut self, node: &mut AVLTreeNode<K, V>, direction: usize) -> &mut AVLTreeNode<K, V> {
        let mut new_root = node.children[1 - direction].take().unwrap();
        node.children[1 - direction] = new_root.children[direction].take();
        new_root.children[direction] = Some(Box::new(mem::take(node)));

        self.node_replace(node, Some(new_root));

        self.update_height(node);
        self.update_height(new_root.as_mut());

        new_root.as_mut()
    }

    /// Balances the subtree rooted at the given node.
    fn node_balance(&mut self, node: &mut AVLTreeNode<K, V>) -> &mut AVLTreeNode<K, V> {
        let left_subtree = &node.children[0];
        let right_subtree = &node.children[1];

        let diff = Self::subtree_height(right_subtree) - Self::subtree_height(left_subtree);

        if diff >= 2 {
            let child = right_subtree.as_ref().unwrap();
            if Self::subtree_height(&child.children[1]) < Self::subtree_height(&child.children[0]) {
                self.rotate(child.as_mut(), 1);
            }
            self.rotate(node, 0)
        } else if diff <= -2 {
            let child = left_subtree.as_ref().unwrap();
            if Self::subtree_height(&child.children[0]) < Self::subtree_height(&child.children[1]) {
                self.rotate(child.as_mut(), 0);
            }
            self.rotate(node, 1)
        } else {
            node
        }
    }

    /// Balances the tree from the given node up to the root.
    fn balance_to_root(&mut self, node: &mut AVLTreeNode<K, V>) {
        let mut rover = Some(node);
        while let Some(node) = rover {
            rover = node.parent.map(|p| unsafe { &mut *p });
            self.node_balance(node);
        }
    }

    /// Inserts a new key-value pair into the tree.
    pub fn insert(&mut self, key: K, value: V) -> bool {
        let mut rover = &mut self.root_node;
        let mut previous_node = None;

        while let Some(ref mut node) = rover {
            previous_node = Some(node.as_mut());
            match (self.compare_func)(&key, &node.key) {
                Ordering::Less => rover = &mut node.children[0],
                Ordering::Greater => rover = &mut node.children[1],
                Ordering::Equal => return false,
            }
        }

        let new_node = Box::new(AVLTreeNode::new(key, value));
        *rover = Some(new_node);

        if let Some(node) = previous_node {
            self.balance_to_root(node);
        }

        self.num_nodes += 1;
        true
    }

    /// Removes a node with the given key from the tree.
    pub fn remove(&mut self, key: &K) -> bool {
        let node = self.lookup_node(key);
        if let Some(node) = node {
            self.remove_node(node);
            true
        } else {
            false
        }
    }

    /// Looks up a node by key.
    fn lookup_node(&self, key: &K) -> Option<&mut AVLTreeNode<K, V>> {
        let mut rover = &mut self.root_node;
        while let Some(ref mut node) = rover {
            match (self.compare_func)(key, &node.key) {
                Ordering::Equal => return Some(node.as_mut()),
                Ordering::Less => rover = &mut node.children[0],
                Ordering::Greater => rover = &mut node.children[1],
            }
        }
        None
    }

    /// Returns the number of nodes in the tree.
    pub fn num_entries(&self) -> usize {
        self.num_nodes
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_lookup() {
        let mut tree = AVLTree::new(|a: &i32, b: &i32| a.cmp(b));
        assert!(tree.insert(10, "ten"));
        assert!(tree.insert(20, "twenty"));
        assert!(tree.insert(5, "five"));

        assert_eq!(tree.lookup_node(&10).unwrap().value, "ten");
        assert_eq!(tree.lookup_node(&20).unwrap().value, "twenty");
        assert_eq!(tree.lookup_node(&5).unwrap().value, "five");
        assert!(tree.lookup_node(&30).is_none());
    }

    #[test]
    fn test_remove() {
        let mut tree = AVLTree::new(|a: &i32, b: &i32| a.cmp(b));
        assert!(tree.insert(10, "ten"));
        assert!(tree.insert(20, "twenty"));
        assert!(tree.insert(5, "five"));

        assert!(tree.remove(&20));
        assert!(tree.lookup_node(&20).is_none());
        assert_eq!(tree.num_entries(), 2);
    }
}
